import sys

import torch
from tqdm import tqdm

from config import Config as C
from model import resnet50
from utils import get_data, get_transform


def main():
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    trans = get_transform()
    _, _, test_set, _, _, test_loader = get_data(
        C.ROOT_PATH,
        C.BATCH_SIZE,
        transform=trans,
        # target_transform=get_target_trans()
    )
    test_num = len(test_set)
    net = resnet50(num_classes=C.CLASSES_NUM)
    net.load_state_dict(torch.load(C.MODEL_WEIGHT_PATH, map_location='cpu'))
    net.to(device)

    net.eval()
    test_bar = tqdm(test_loader, file=sys.stdout)
    acc = 0.0  # accumulate accurate number / epoch
    with torch.no_grad():
        for val_data in test_bar:
            val_images, val_labels = val_data
            outputs = net(val_images.to(device))
            # loss = loss_function(outputs, test_labels)
            predict_y = torch.max(outputs, dim=1)[1]
            acc += torch.eq(predict_y, val_labels.to(device)).sum().item()
    print(f"Accuracy is {(acc / test_num):.3%}")


if __name__ == '__main__':
    main()
